package themis.flink.riskengine;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import java.util.List;
import org.apache.commons.lang3.StringUtils;
import org.apache.flink.api.common.functions.RichMapFunction;

public class FeatureEnrichFunction extends RichMapFunction<JSONObject, JSONObject> {
    private List<Rule> features;

    public FeatureEnrichFunction(List<Rule> features) {
        this.features = features;
    }

    public JSONObject map(JSONObject value) throws Exception {
        String keyName = value.getString("KEY_NAME");
        if (keyName.equals("event"))
            return value;
        for (Rule feature : this.features) {
            Object featureResult;
            String key = feature.getMeasure() + "#" + feature.getDimension();
            if (!StringUtils.equals(key, keyName))
                continue;
            String function = feature.getKey();
            long window = feature.getInterval();
            JSONArray history = value.getJSONArray("HISTORY");
            String target = key.replace(".history", "").split("#")[0];
            if ("sum".equalsIgnoreCase(function)) {
                featureResult = doSum(history, target, window);
            } else if ("count".equalsIgnoreCase(function)) {
                featureResult = doCount(history, target, window);
            } else if ("count_distinct".equalsIgnoreCase(function)) {
                featureResult = doCountDistinct(history, target, window);
            } else {
                throw new UnsupportedOperationException(String.format("unsupported function[%s]", function));
            }
            value.putIfAbsent("features", new JSONObject());
            String featureName = String.format("%s(%s,%s)", function, key, window);
            value.getJSONObject("features").put(featureName, featureResult);
        }
        return value;
    }

    private double doSum(JSONArray history, String target, long window) {
        long maxTimestamp = (Long) history.stream()
                .map(x -> ((JSONObject) x).getLong("timestamp"))
                .max(Long::compare)
                .orElse(Long.MIN_VALUE);
        long minTimestamp = maxTimestamp - window;
        return (Double) history.stream()
                .filter(x -> (((JSONObject) x).getLong("timestamp") > minTimestamp))
                .map(x -> ((JSONObject) x).getDouble(target))
                .reduce(Double::sum)
                .orElse(Double.MIN_VALUE);
    }

    private double doCount(JSONArray history, String target, long window) {
        long maxTimestamp = (Long) history.stream()
                .map(x -> ((JSONObject) x).getLong("timestamp"))
                .max(Long::compare)
                .orElse(Long.MIN_VALUE);
        long minTimestamp = maxTimestamp - window;
        return history.stream()
                .filter(x -> (((JSONObject) x).getLong("timestamp") > minTimestamp))
                .count();
    }

    private double doCountDistinct(JSONArray history, String target, long window) {
        long maxTimestamp = (Long) history.stream()
                .map(x -> ((JSONObject) x).getLong("timestamp"))
                .max(Long::compare)
                .orElse(Long.MIN_VALUE);
        long minTimestamp = maxTimestamp - window;
        return history.stream()
                .filter(x -> (((JSONObject) x).getLong("timestamp") > minTimestamp))
                .map(x -> ((JSONObject)x).getString(target))
                .distinct()
                .count();
    }
}
